#  udp.rb
#
#  This file is part of myvnet.
#
#  Copyright (C) 2008  Filipe David Borba Manana,  fdmanana@gmail.com
#
#  myvnet is free software: you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#
#  myvnet is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with myvnet.  If not, see <http://www.gnu.org/licenses/>.


require 'lib/protocols/generic_protocol'
require 'lib/protocols/ipv4'


# User Datagram Protocol (UDP) module.
#
module UDP
        
  # UDP packet.
  # 
  class Packet < GenericProtocol::Packet
    
    HLEN = 8
    
    @packet_type = :udp
    @layer = 4
    
    attr_accessor :sport     # source port
    attr_accessor :dport     # destination port
    attr_accessor :len       # length (in bytes)
    attr_accessor :checksum  # covers header and data (payload)
    
    
    def initialize(bstream = nil)
      super
      
      bstream = ("\x00" * 5 + "\x08" + "\x00" * 2) if bstream.nil?
      
      @sport = (bstream[0,2].unpack 'n')[0]
      @dport = (bstream[2,2].unpack 'n')[0]
      @len = (bstream[4,2].unpack 'n')[0]
      @checksum = (bstream[6,2].unpack 'n')[0]
      
      plen = @len - HLEN
      @payload = bstream[HLEN, plen]
    end
        

    # === Synopsis
    #
    #   UDP::Packet#to_bstream(pseudo_header = '')
    #   
    # === Return
    #
    # A byte stream, in network byte order, representation of the receiver
    # UDP packet.
    #
    # === Description
    #
    # Builds a byte stream (String instance), in network byte order, 
    # representation of the receiver UDP packet.
    #
    # *Notes*: 
    # * It calculates and sets the length field for the receiving UDP packet.
    # * If the receiving packet's checksum is zero, the checksum is calculated
    #   and set.
    #            
    def to_bstream
      stream = ''      
      stream << ( [ @sport ].pack 'n')
      stream << ( [ @dport ].pack 'n')
      stream << "\x00\x00"
      stream << ( [ @checksum ].pack 'n')

      payload_stream = ''
      
      case @payload
        when String
          payload_stream << @payload
        when GenericProtocol::Packet
          payload_stream << (@payload.to_bstream)
      end
      
      @len = HLEN + payload_stream.length
      stream[4,2] = ( [@len].pack 'n')
      
      stream << payload_stream
      
      if @checksum.zero?
        
        pseudo_header = ''
        
        if @parent.is_a? IPv4::Packet
          # defined in RFC 768
          pseudo_header << IPv4.to_bstream_addr(@parent.src)
          pseudo_header << IPv4.to_bstream_addr(@parent.dst)
          pseudo_header << ("\x00\x11" + stream[4,2])
        end
        
        set_checksum(pseudo_header, stream)
      end
      
      stream
    end
    
    
    def build_upper_layers!
      # no support for application layer protocols, yet
      self
    end
    
    
    # === Synopsis
    # 
    #   UDP::Packet#len_calc!
    #   
    # === Return
    # 
    # The computed length value for the packet's length field.
    # 
    # === Description
    # 
    # Computes the length of a UDP packet and sets the receiver
    # packet's length field to the computed value.
    # 
    def len_calc!
      @len = self.byte_length
    end 

    
    # === Synopsis
    # 
    #   UDP::Packet#checksum_calc!(pseudo_header = '')
    #
    # === Return
    # 
    # The checksum that was computed for the UDP packet.
    # 
    # === Description
    # 
    # Computes the checksum of a UDP packet and sets the receiver
    # packet's checksum field value to the computed value.
    #    
    def checksum_calc!
      @checksum = 0
      self.to_bstream
      @checksum
    end
    
    
    def to_s(indent_level = 0)
      prefix = ' ' * indent_level
      str = ''
      str << prefix << "#{self.class.name}\n"
      str << prefix << "src_port: #{@sport}, dst_port: #{@dport}, "
      str << "len: #{@len}, chksum: 0x" << ("%04X" % @checksum) << ",\n"
      str << prefix << "payload:\n"
      
      case @payload
        when String
          str << prefix << '   ' << hexdump(@payload) << "\n"
        when GenericProtocol::Packet
          str << @payload.to_s(indent_level + 3)
      end
      
      str
    end    


    private
    
    def set_checksum(pseudo_header, bstream)  
      tmpstr = pseudo_header + bstream
      
      # if odd number of bytes, add an extra one
      tmpstr << "\x00" if tmpstr.length[0] == 1
      
      words = tmpstr.unpack 'n*'
      
      sum = words.inject do |w,a|
        w + a
      end

      sum = (sum & 0xff_ff) + (sum >> 16)
      chksum = (~sum) & 0xff_ff
      chksum = 0xff_ff if chksum.zero?      
      
      bstream[6,2] = [chksum].pack 'n'
      @checksum = chksum

      bstream
    end    
                
  end
  
end